package com.LeeCode;

import java.util.*;

/**
 * 二叉树中的第 K 大层和
 */

public class Code2583 {
    public static void main(String[] args) {
        Integer[] arr = {5, 8, 9, 2, 1, 3, 7, 4, 6};
        TreeNode root = Utils.buildTree(arr);
        System.out.println(new Code2583().kthLargestLevelSum(root, 2));
    }

    public long kthLargestLevelSum(TreeNode root, int k) {
        List<Long> ans = new ArrayList<>(k);
        Deque<TreeNode> queue = new LinkedList<>();
        queue.offer(root);
        while (!queue.isEmpty()) {
            int size = queue.size();
            long sum = 0;
            for (int i = 0; i < size; i++) {
                TreeNode node = queue.poll();
                sum += node.val;
                if (node.left != null) queue.offer(node.left);
                if (node.right != null) queue.offer(node.right);
            }
            ans.add(sum);
        }
        Collections.sort(ans, Collections.reverseOrder());
        if (k > ans.size()) return -1;
        return ans.get(k - 1);
    }

    public long kthLargestLevelSum1(TreeNode root, int k) {
        // 小顶堆
        PriorityQueue<Long> pq = new PriorityQueue<>(k);
        // BFS
        Queue<TreeNode> que = new LinkedList<>();
        que.offer(root);
        while (!que.isEmpty()) {
            long sum = 0L;
            for (int i = que.size(); i > 0; i--) {
                TreeNode cur = que.poll();
                sum += cur.val;
                if (cur.left != null) {
                    que.offer(cur.left);
                }
                if (cur.right != null) {
                    que.offer(cur.right);
                }
            }
            if (pq.size() < k || sum > pq.peek()) {
                if (pq.size() == k) {
                    pq.poll();
                }
                pq.offer(sum);
            }
        }
        return pq.size() == k ? pq.peek() : -1;
    }
}
